functor ConstFold (S : CONST_FOLD_STRUCTS) : CONST_FOLD =
struct
  open S
  
  structure I = Ir
  structure ErrorMsg = ErrorMsg
  
  local open I
  in structure Id = I.Id
     structure Intl = I.Intl
     structure B = I.Binop
     structure T = I.Types
     structure L = I.Le
     structure S = I.Structs
     structure M = I.Stm
     structure F = I.Function
     structure P = I.Program
  end

  val error = ErrorMsg.error

  val tag = false
  fun debugPrint s = 
      if tag
      then print (String.concat [s, "\n"])
      else ()

  val doAgain = ref false
  fun repeat () = doAgain := true
  fun reset () = doAgain := false
  
  fun shift b =
      if b
      then L.True
      else L.False
  
  fun cfOneExp exp =
      case exp
       of L.Expvoid => L.Expvoid
        | L.IntNum i => L.IntNum i
        | L.True => L.True
        | L.False => L.False
        | L.Null => L.Null
        | L.Lval {lval, ty} => L.Lval {lval = cfLval lval, ty = ty}
        | L.BinopExp {left = L.IntNum i, bop = B.Plus, right = L.IntNum j, ty, pos} =>
          let val _ = repeat ()
              val new = L.IntNum (Intl.plus (i, j))
          in  new
          end
        | L.BinopExp {left, bop = B.Plus, right, ty, pos} =>
          L.BinopExp {left = cfOneExp left,
                      bop = B.Plus,
                      right = cfOneExp right,
                      ty = ty,
                      pos = pos}
        | L.BinopExp {left = L.IntNum i, bop = B.Minus, right = L.IntNum j, ty, pos} =>
          let val _ = repeat ()
              val new = L.IntNum (Intl.minus (i, j))
          in  new
          end
        | L.BinopExp {left, bop = B.Minus, right, ty, pos} =>
          L.BinopExp {left = cfOneExp left,
                      bop = B.Minus,
                      right = cfOneExp right,
                      ty = ty,
                      pos = pos}
        | x as L.BinopExp {left = L.IntNum i, bop = B.Times, right, ty, pos} =>
          if Intl.isZero i
          then let val _ = repeat ()
                   val new = L.IntNum (Intl.zero)
               in  new
               end
          else (case right
                of L.IntNum j =>
                   let val _ = repeat ()
                       val new = L.IntNum (Intl.times (i, j))
                   in  new 
                   end
                 | _ => x)
        | x as L.BinopExp {left = e1, bop = B.Times, right = L.IntNum j, ty, pos} =>
          if Intl.isZero j
          then let val _ = repeat ()
                   val new = L.IntNum (Intl.zero)
               in  new
               end
          else (case e1
                 of L.IntNum i =>
                    let val _ = repeat ()
                        val new = L.IntNum (Intl.times (i,j))
                    in  new
                    end
                  | _ => x)
        | L.BinopExp {left, bop = B.Times, right, ty, pos} =>
          L.BinopExp {left = cfOneExp left,
                      bop = B.Times,
                      right = cfOneExp right,
                      ty = ty,
                      pos = pos}
        | x as L.BinopExp {left = L.IntNum i, bop = B.Divide, right, ty, pos} =>
          if Intl.isZero i
          then let val _ = repeat ()
                   val new = L.IntNum (Intl.zero)
               in  new
               end
          else (case right
                of L.IntNum j =>
                   if Intl.isZero j
                   then (error pos "divide by zero"
                       ; L.IntNum (Intl.zero))
                   else let val _ = repeat ()
                            val new = L.IntNum (Intl.divide (i, j))
                        in  new
                        end
                 | _ => x)
        | L.BinopExp {left, bop = B.Divide, right, ty, pos} =>
          L.BinopExp {left = cfOneExp left,
                      bop = B.Divide,
                      right = cfOneExp right,
                      ty = ty,
                      pos = pos}
        | x as L.BinopExp {left = L.IntNum i, bop = B.Modus, right, ty, pos} =>
          if Intl.isZero i
          then let val _ = repeat ()
                   val new = L.IntNum (Intl.zero)
               in  new
               end
          else (case right
                of L.IntNum j =>
                   if Intl.isZero j
                   then (error pos "divided by zero"
                       ; L.IntNum (Intl.zero))
                   else let val _ = repeat ()
                            val new = L.IntNum (Intl.modus (i, j))
                        in  new
                        end
                 | _ => x)
        | L.BinopExp {left, bop = B.Modus, right, ty, pos} => 
          L.BinopExp {left = cfOneExp left,
                      bop = B.Modus,
                      right = cfOneExp right,
                      ty = ty,
                      pos = pos}
        | L.BinopExp {left = L.True, bop = B.And, right, ty, pos} =>
          let val _ = repeat ()
          in  right
          end
        | L.BinopExp {left = L.False, bop = B.And, right, ty, pos} =>
          let val _ = repeat ()
          in  L.False
          end
        | L.BinopExp {left, bop = B.And, right, ty, pos} =>
          L.BinopExp {left = cfOneExp left,
                      bop = B.And,
                      right = cfOneExp right,
                      ty = ty,
                      pos = pos}
        | L.BinopExp {left = L.True, bop = B.Or, right, ty, pos} =>
          let val _ = repeat ()
          in  L.True
          end
        | L.BinopExp {left = L.False, bop = B.Or, right, ty, pos} =>
          let val _ = repeat ()
          in  right
          end
        | L.BinopExp {left, bop = B.Or, right, ty, pos} => 
          L.BinopExp {left = cfOneExp left,
                      bop = B.Or,
                      right = cfOneExp right,
                      ty = ty,
                      pos = pos}
        | L.BinopExp {left = L.IntNum i, bop = B.Eq, right = L.IntNum j, ty, pos} =>
          let val _ = repeat ()
          in  shift (Intl.equals (i, j))
          end
        | L.BinopExp {left, bop = B.Eq, right, ty, pos} =>
          L.BinopExp {left = cfOneExp left,
                      bop = B.Eq,
                      right = cfOneExp right,
                      ty = ty, 
                      pos = pos}
        | L.BinopExp {left = L.IntNum i, bop = B.Ne, right = L.IntNum j, ty, pos} =>
          let val _ = repeat ()
          in  shift (Intl.ne (i, j))
          end
        | L.BinopExp {left, bop = B.Ne, right, ty, pos} =>
          L.BinopExp {left = cfOneExp left,
                      bop = B.Ne,
                      right = cfOneExp right,
                      ty = ty,
                      pos = pos}
        | L.BinopExp {left = L.IntNum i, bop = B.Gt, right = L.IntNum j, ty, pos} =>
          let val _ = repeat ()
          in  shift (Intl.gt (i, j))
          end
        | L.BinopExp {left, bop = B.Gt, right, ty, pos} =>
          L.BinopExp {left = cfOneExp left,
                      bop = B.Gt,
                      right = cfOneExp right,
                      ty = ty,
                      pos = pos}
        | L.BinopExp {left = L.IntNum i, bop = B.Ge, right = L.IntNum j, ty, pos} =>
          let val _ = repeat ()
          in  shift (Intl.ge (i, j))
          end
        | L.BinopExp {left, bop = B.Ge, right, ty, pos} =>
          L.BinopExp {left = cfOneExp left,
                      bop = B.Ge,
                      right = cfOneExp right,
                      ty = ty,
                      pos = pos}
        | L.BinopExp {left = L.IntNum i, bop = B.Lt, right = L.IntNum j, ty, pos} =>
          let val _ = repeat ()
          in  shift (Intl.lt (i, j))
          end
        | L.BinopExp {left, bop = B.Lt, right, ty, pos} =>
          L.BinopExp {left = cfOneExp left,
                      bop = B.Lt,
                      right = cfOneExp right,
                      ty = ty,
                      pos = pos}
        | L.BinopExp {left = L.IntNum i, bop = B.Le, right = L.IntNum j, ty, pos} =>
          let val _ = repeat ()
          in  shift (Intl.le (i, j))
          end
        | L.BinopExp {left, bop = B.Le, right, ty, pos} =>
          L.BinopExp {left = cfOneExp left,
                      bop = B.Le,
                      right = cfOneExp right,
                      ty = ty,
                      pos = pos}
        | L.NotExp (L.True, pos) =>
          let val _ = repeat ()
          in  L.False
          end
        | L.NotExp (L.False, pos) =>
          let val _ = repeat ()
          in  L.True
          end
        | L.NotExp (e, pos) =>
          L.NotExp (cfOneExp e, pos)
        | L.Result => L.Result
        | L.Imply (L.True, e2, pos) =>
          let val _ = repeat ()
          in  e2
          end
        | L.Imply (L.False, e2, pos) =>
          let val _ = repeat ()
          in  L.True
          end
        | L.Imply (e1, L.True, pos) =>
          let val _ = repeat ()
          in  L.True
          end
        | L.Imply (e1, L.False, pos) =>
          let val _ = repeat ()
          in  L.NotExp (e1, pos)
          end
        | L.Imply (e1, e2, pos) => L.Imply (cfOneExp e1, cfOneExp e2, pos)
        | L.Iff (e1, e2, pos) =>
          let val _ = repeat ()
          in  L.BinopExp {left = L.Imply (e1, e2, pos),
                      bop = B.And,
                      right = L.Imply (e2, e1, pos),
                      ty = T.Boolt,
                      pos = pos}
          end
        | L.Forall (id, ty, e, pos) => L.Forall (id, ty, e, pos)
        | L.Exists (id, ty, e, pos) => L.Exists (id, ty, e, pos)
        | L.Pi (lss, pos) => L.Pi (lss, pos)
        | L.N (ls, pos) => L.N (ls, pos)
        | L.D (ls, pos) => L.D (ls, pos)
        | L.Prev (lval, pos) => L.Prev (lval, pos)
        | L.LList (exp, pos) => L.LList (exp, pos)
        | L.Tree (exp, pos) => L.Tree (exp, pos)
        | L.Pred (id, exp, pos) => L.Pred (id, exp, pos)
        | L.Let (e1, e2, e3, pos) => L.Let (e1, e2, e3, pos)
      
  and cfExp exp =
      let fun loop e =
              let val res = cfOneExp e
              in  if L.equalsExp (res, e)
                  then res
                  else (debugPrint "loop in exp optimization"
                      ; loop res)
              end
      in  loop exp
      end
      
  and cfLval lval = lval
  
  and mkSomeList x = SOME [x]
  
  and mkSome x = SOME x
  
  and cfStm (stm : M.t) : M.t list option =
      case stm
       of M.AsnExp {lval, exp, ty, pos} =>
          let val new = M.AsnExp {lval = cfLval lval,
                                  exp = cfExp exp,
                                  ty = ty,
                                  pos = pos}
          in  mkSomeList new
          end 
        | M.AsnCall {lval, fid, exps, ty, pos} =>
          let val new = M.AsnCall {lval = lval,
                                   fid = fid,
                                   exps = List.map cfExp exps,
                                   ty = ty,
                                   pos = pos}
          in  mkSomeList new
          end
        | M.AsnAlloc {lval, ty, pos} =>
          let val new = M.AsnAlloc {lval = cfLval lval,
                                    ty = ty,
                                    pos = pos}
          in  mkSomeList new
          end
        | M.If (L.True, s1, s2, pos) => mkSome s1
        | M.If (L.False, s1, s2, pos) => mkSome s2
        | M.If (e, s1, s2, pos) =>
          let val new = M.If (cfExp e, cfStms s1, cfStms s2, pos)
          in  mkSomeList new
          end
        | M.While {condition = L.False, invariant, body, pos} => NONE
        | M.While {condition = L.True, invariant, body, pos} =>
          let val new = M.While {invariant = invariant,
                                 condition = L.True,
                                 body = cfStms body,
                                 pos = pos}
          in  mkSomeList new
          end
        | M.While {invariant, condition, body, pos} =>
          let val new = M.While {invariant = invariant,
                                 condition = cfExp condition,
                                 body = cfStms body,
                                 pos = pos}
          in  debugPrint "loop in while"
            ; mkSomeList new
          end
        | M.Return {exp, fty, pos} =>
          let val new = M.Return {exp = cfExp exp,
                                  fty = fty,
                                  pos = pos}
          in  mkSomeList new
          end
        | M.Call (id, es, pos) =>
          let val new = M.Call (id, List.map cfExp es, pos) 
          in  mkSomeList new
          end
        | M.Free (e, pos) =>
          let val new = M.Free (cfExp e, pos)
          in  mkSomeList new
          end
        | M.Prop (e, pos) =>
          let val new = M.Prop (cfExp e, pos)
          in  mkSomeList new
          end
  
  and clean (l : M.t list option list) : M.t list =
      let fun f l =
              case l
               of [] => []
                | (NONE :: xs) => 
                  let val _ = repeat ()
                  in  f xs
                  end
                | (SOME x :: xs) => x :: (f xs)
      in  List.concat (f l)
      end
      
  and cfStms (ss : M.t list) : M.t list = clean (List.map cfStm ss)
     
  fun cfTopStms stms = 
      let val stms = cfStms stms
      in  if !doAgain
          then (reset (); cfTopStms stms)
          else stms
      end
      
  fun cfFunction f =
      case f
       of F.T {name,
               rettype,
               args,
               locals,
               stm,
               pre,
               post,
               pos} => 
          F.T {name = name,
               rettype = rettype,
               args = args,
               locals = locals,
               stm = cfTopStms stm,
               pre = pre,
               post = post,
               pos = pos}
      
  fun shrink p =
      case p
       of P.T {structs,
               gvars,
               funcs,
               pos} =>
          P.T {structs = structs,
               gvars = gvars,
               funcs = List.map cfFunction funcs,
               pos = pos} 
end
